const Koa = require('koa')
const Router = require('koa-router')
const json = require('koa-json')
const bodyParser = require('koa-bodyparser')
const cors = require('@koa/cors')
const fs = require('fs')
const llm = require('./llm')
const models = require('./models.json')

const app = new Koa()
app.use(cors())
app.use(require('koa-static')(__dirname + '/public', { maxAge : 7 * 24 * 60 * 60 * 1000 }))
app.use(bodyParser({enableTypes:['json', 'form', 'text']}))
app.use(json())

const router = new Router()

router.get('/v1/models', async (ctx) => {
  ctx.body = {
    "object": "list",
    "data": models.map(m=>{
      return { id: m.filename, name: m.filename, exist: fs.existsSync(`./models/${m.filename}`) }
    }) 
  }
})

router.post('/v1/change_model', async (ctx) => {
  await llm.load(ctx.request.body.model)
  ctx.body = {}
})

router.get('/c/:id', async(ctx)=>{
  ctx.type = 'html'
  ctx.body = fs.createReadStream(__dirname + '/public/index.html')
})

router.post('/v1/chat/completions', async (ctx) => {
  ctx.set('Content-Type', 'text/event-stream');
  ctx.set('Cache-Control', 'no-cache');
  ctx.set('Connection', 'keep-alive');
  ctx.request.body.messages.forEach(element => {
    element.content = element.content[0].text
  })
  if(ctx.request.body.id && ctx.request.body.id > 0){
    ctx.body = await llm.complete(ctx.request.body.id, ctx.request.body.model, ctx.request.body.messages)
  }else{
    ctx.body = await llm.generate(ctx.request.body.model, ctx.request.body.messages)
  }
})

app.use(router.routes())
app.use(router.allowedMethods())

app.listen(1337, () => {
  console.log(`Server is running at http://localhost:1337`);
})
